//go:build windows
// +build windows

package process

import (
	"fmt"
	"go-cmd/internal/common"
	"go-cmd/top/cpu"
	"go-cmd/top/mem"
	"unsafe"

	"golang.org/x/sys/windows"
)

// load Windows API functions
// You can get more information from https://learn.microsoft.com/zh-cn/windows/win32/api/tlhelp32/
// and https://learn.microsoft.com/zh-CN/windows/win32/api/handleapi/
// Such document can give you the example of how to call this function
var (
	modKernel32                  = common.ModKernel32
	procCloseHandle              = modKernel32.NewProc("CloseHandle")
	procCreateToolHelp32Snapshot = modKernel32.NewProc("CreateToolhelp32Snapshot")
	procProcess32First           = modKernel32.NewProc("Process32FirstW")
	procProcess32Next            = modKernel32.NewProc("Process32NextW")

	// this is used to get access
	// You can see https://learn.microsoft.com/zh-cn/windows/win32/api/winbase/nf-winbase-lookupprivilegevaluea
	// to get more information.
	advapi32                  = common.ModAdvapi32
	procLookupPrivilegeValue  = advapi32.NewProc("LookupPrivilegeValueW")
	procAdjustTokenPrivileges = advapi32.NewProc("AdjustTokenPrivileges")
)

const (
	ERROR_NO_MORE_FILES = 0x12
	MAX_PATH            = 260
)

type winLUID struct {
	LowPart  uint32
	HighPart int32
}

type winLUIDAndAttributes struct {
	Luid       winLUID
	Attributes uint32
}

type winTokenPrivileges struct {
	PrivilegeCount uint32
	Privileges     [1]winLUIDAndAttributes
}

// We need to get the privilege to get information of process.
// U can get the example of how to solve by C++ at:
// https://learn.microsoft.com/zh-cn/windows/win32/secauthz/enabling-and-disabling-privileges-in-c--
// But it seems like that the method is not usable
func init() {
	handle := windows.CurrentProcess()
	var token windows.Token
	err := windows.OpenProcessToken(handle, windows.TOKEN_ALL_ACCESS, &token)
	if err != nil {
		fmt.Printf("Get Access to System failed, err=%s\n", err)
		return
	}
	defer token.Close()

	tokenPrivileges := winTokenPrivileges{
		PrivilegeCount: 1,
	}

	// SE_DEBUG_NAME TEXT (“SeDebugPrivilege”)
	// Required to debug and adjust the memory of a process owned by another account.
	// User Right: Debug programs.
	// U can also see other info at https://learn.microsoft.com/en-us/windows/win32/secauthz/privilege-constants
	lpName, err := windows.UTF16FromString("SeDebugPrivilege")
	if err != nil {
		fmt.Println("Get SE_DEBUG_NAME Failed, err-", err.Error())
		return
	}
	ret, _, err := procLookupPrivilegeValue.Call(
		0,
		uintptr(unsafe.Pointer(&lpName[0])),
		uintptr(unsafe.Pointer(&tokenPrivileges.Privileges[0].Luid)),
	)
	if ret == 0 && err != nil {
		fmt.Printf("Lookup Privilege failed, err=%s\n", err.Error())
		return
	}
	tokenPrivileges.Privileges[0].Attributes = windows.SE_PRIVILEGE_ENABLED

	ret, _, err = procAdjustTokenPrivileges.Call(
		uintptr(token),
		0,
		uintptr(unsafe.Pointer(&tokenPrivileges)),
		uintptr(unsafe.Sizeof(tokenPrivileges)),
		0,
		0,
	)
	if ret == 0 && err != nil {
		fmt.Printf("Adjust Token Privileges Failed, err=%s\n", err.Error())
		return
	}
}

// PROCESSENTRY32 is the Windows API structure that contains a process's information
// You can read https://learn.microsoft.com/zh-cn/windows/win32/api/tlhelp32/ns-tlhelp32-processentry32
// for more information
type PROCESSENTRY32 struct {
	Size              uint32
	CntUsage          uint32
	ProcessID         uint32
	DefaultHeapID     uintptr
	ModuleID          uint32
	CntThreads        uint32
	ParentProcessID   uint32
	PriorityClassBase int32
	Flags             uint32
	ExeFile           [MAX_PATH]uint16
}

// the implementation of Process which can run on windows
type WindowsProcess struct {
	pid  int
	ppid int
	exe  string
}

func (p *WindowsProcess) Pid() int {
	return p.pid
}

func (p *WindowsProcess) PPid() int {
	return p.ppid
}

func (p *WindowsProcess) Execute() string {
	return p.exe
}

func (p *WindowsProcess) Memory() (uint64, error) {
	mem, err := mem.GetMemoryInfo(p.Pid())
	if err != nil {
		return 0, err
	}
	return mem.RSS, nil
}

func (p *WindowsProcess) Cpu() (uint64, error) {
	ret, err := cpu.GetProcessCPUTime(p.pid)
	if err != nil {
		return 0, err
	}
	return ret.Kernel + ret.User, nil
}

func newWindowsProcess(e *PROCESSENTRY32) *WindowsProcess {
	end := 0
	for {
		if e.ExeFile[end] == 0 {
			break
		}
		end++
	}

	return &WindowsProcess{
		pid:  int(e.ProcessID),
		ppid: int(e.ParentProcessID),
		exe:  windows.UTF16ToString(e.ExeFile[:end]),
	}
}

func findProcess(pid int) (Process, error) {
	ps, err := processes()
	if err != nil {
		return nil, err
	}
	for _, p := range ps {
		if p.Pid() == pid {
			return p, nil
		}
	}

	return nil, nil
}

func processes() ([]Process, error) {
	handle, _, _ := procCreateToolHelp32Snapshot.Call(0x00000002, 0)
	if handle < 0 {
		return nil, windows.GetLastError()
	}
	defer procCloseHandle.Call(handle)

	var entry PROCESSENTRY32
	entry.Size = uint32(unsafe.Sizeof(entry))
	ret, _, _ := procProcess32First.Call(handle, uintptr(unsafe.Pointer(&entry)))
	if ret == 0 {
		return nil, fmt.Errorf("Error retrieving process info.")
	}

	results := make([]Process, 0, 50)
	for {
		results = append(results, newWindowsProcess(&entry))

		ret, _, _ := procProcess32Next.Call(handle, uintptr(unsafe.Pointer(&entry)))
		if ret == 0 {
			break
		}
	}
	return results, nil
}
